import numpy as np
from scipy.linalg import block_diag

def F4_MPC_Matrices_PM(A, B, Q, R, S, N_P):
    # 计算系统矩阵维度，n
    n = A.shape[0]
    
    # 计算输入矩阵维度，p
    p = B.shape[1]

    # 初始化Phi矩阵并定义维度
    Phi = np.zeros((N_P * n, n))

    # 初始化Gamma矩阵并定义维度
    Gamma = np.zeros((N_P * n, N_P * p))

    # 定义临时对角单位矩阵
    tmp = np.eye(n)  # Create a n x n identity matrix

    # 定义行向量
    rows = np.arange(n)

    # for循环，用于构建Phi和Gamma矩阵
    for i in range(1, N_P + 1):
        # 构建Phi矩阵，参考式（5.3.5b）
        Phi[(i - 1) * n:i * n, :] = np.linalg.matrix_power(A, i)
        
        # 构建Gamma矩阵，参考式（5.3.5b）
        Gamma[rows, :] = np.hstack((tmp @ B, Gamma[max(0, rows[0] - n):(max(0, rows[0] - n) + n), :N_P * p - p]))
        
        # 更新行数
        rows = i * n + np.arange(n)
        
        # 更新临时矩阵用于矩阵A的幂计算
        tmp = A @ tmp

    # 构建Omega矩阵，包含Q矩阵的部分
    Omega = np.kron(np.eye(N_P - 1), Q)

    # 构建最终Omega矩阵，包含S矩阵
    Omega = block_diag(Omega, S)

    # 构建Psi矩阵，其为R矩阵组成的对角阵
    Psi = np.kron(np.eye(N_P), R)

    # 计算二次规划矩阵F
    F = Gamma.T @ Omega @ Phi

    # 计算二次规划矩阵H
    H = Psi + Gamma.T @ Omega @ Gamma

    return Phi, Gamma, Omega, Psi, F, H